#!/usr/bin/env python3
"""
MobileViT 树莓派推理程序
用于导盲系统的实时障碍检测
"""

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import cv2
import numpy as np
import time
import threading
import queue
from pathlib import Path
import logging
from typing import Dict, List, Tuple, Optional
import argparse

# 导入MobileViT模型定义
import sys
sys.path.append('..')
from mobilevit_complete import MobileViT

# 导入摄像头和语音模块
from camera_capture import CameraCapture
from voice_alert import VoiceAlert

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('inference.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class UltrasonicReader:
    def __init__(self, trig_pin: int = 10, echo_pin: int = 11):
        self.trig_pin = trig_pin
        self.echo_pin = echo_pin
        self.latest_m = None
        self.running = False
        self.thread = None
        self._gpio = None
        try:
            import OPi.GPIO as GPIO
            self._gpio = GPIO
            self._gpio.setmode(self._gpio.BOARD)
            self._gpio.setup(self.trig_pin, self._gpio.OUT)
            self._gpio.setup(self.echo_pin, self._gpio.IN)
            self._gpio.output(self.trig_pin, self._gpio.LOW)
        except Exception:
            self._gpio = None

    def start(self):
        if not self._gpio or self.running:
            return False
        self.running = True
        self.thread = threading.Thread(target=self._loop, daemon=True)
        self.thread.start()
        return True

    def _loop(self):
        interval = 0.2
        samples = []
        max_samples = 5
        while self.running:
            try:
                g = self._gpio
                g.output(self.trig_pin, g.HIGH)
                time.sleep(0.00001)
                g.output(self.trig_pin, g.LOW)

                t0 = time.perf_counter()
                timeout = 0.03
                while g.input(self.echo_pin) == g.LOW:
                    if time.perf_counter() - t0 > timeout:
                        raise TimeoutError
                start = time.perf_counter()
                while g.input(self.echo_pin) == g.HIGH:
                    if time.perf_counter() - start > timeout:
                        break
                end = time.perf_counter()
                pulse = end - start
                dist_m = (pulse * 343.0) / 2.0
                if 0.05 <= dist_m <= 4.0:
                    samples.append(dist_m)
                    if len(samples) > max_samples:
                        samples.pop(0)
                    med = sorted(samples)[len(samples)//2]
                    self.latest_m = med
                time.sleep(interval)
            except Exception:
                time.sleep(interval)

    def stop(self):
        if not self._gpio:
            return
        self.running = False
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=2)
        try:
            self._gpio.cleanup()
        except Exception:
            pass

    def get_distance_m(self) -> Optional[float]:
        return self.latest_m

class ObstacleDetector:
    """障碍检测器主类"""
    
    def __init__(self, model_path: str, device: str = 'cpu'):
        """
        初始化障碍检测器
        
        Args:
            model_path: 模型权重文件路径
            device: 推理设备 ('cpu' 或 'cuda')
        """
        self.device = torch.device(device)
        self.model_path = model_path
        
        # 类别映射
        self.class_names = ['台阶', '杆子', '正常', '路口', '路障']
        self.risk_levels = {
            '正常': 'low',
            '台阶': 'medium', 
            '杆子': 'medium',
            '路口': 'high',
            '路障': 'high'
        }
        
        # 初始化模型
        self.model = self._load_model()
        
        # 图像预处理
        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])
        
        # 初始化摄像头和语音模块
        self.camera = CameraCapture(resolution=(320, 240))
        self.voice_alert = VoiceAlert()
        self.ultrasonic = UltrasonicReader(10, 11)
        self.ultrasonic.start()
        
        # 推理队列和线程控制
        self.frame_queue = queue.Queue(maxsize=5)
        self.result_queue = queue.Queue(maxsize=10)
        self.running = False
        
        # 性能统计
        self.inference_times = []
        self.fps_counter = 0
        self.last_fps_time = time.time()
        
        logger.info("障碍检测器初始化完成")
    
    def _load_model(self) -> nn.Module:
        """加载MobileViT模型"""
        try:
            # 创建模型实例
            model = MobileViT(
                image_size=(224, 224),
                dims=[144, 192, 240],
                channels=[16, 32, 64, 64, 96, 96, 128, 128, 160, 160, 640],
                num_classes=5
            )
            
            # 加载权重
            checkpoint = torch.load(self.model_path, map_location=self.device)
            model.load_state_dict(checkpoint)
            model.to(self.device)
            model.eval()
            
            logger.info(f"模型加载成功: {self.model_path}")
            return model
            
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            raise
    
    def preprocess_frame(self, frame: np.ndarray) -> torch.Tensor:
        """预处理图像帧"""
        # 转换BGR到RGB
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # 应用变换
        tensor = self.transform(frame_rgb)
        tensor = tensor.unsqueeze(0)  # 添加batch维度
        
        return tensor.to(self.device)
    
    def predict(self, frame: np.ndarray) -> Dict:
        """对单帧进行预测"""
        start_time = time.time()
        
        try:
            # 预处理
            input_tensor = self.preprocess_frame(frame)
            
            # 推理
            with torch.no_grad():
                outputs = self.model(input_tensor)
                probabilities = torch.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probabilities, 1)
            
            # 解析结果
            class_idx = predicted.item()
            class_name = self.class_names[class_idx]
            confidence_score = confidence.item()
            risk_level = self.risk_levels[class_name]
            
            inference_time = time.time() - start_time
            self.inference_times.append(inference_time)
            
            # 保持最近100次推理时间
            if len(self.inference_times) > 100:
                self.inference_times.pop(0)
            
            result = {
                'class_name': class_name,
                'confidence': confidence_score,
                'risk_level': risk_level,
                'inference_time': inference_time,
                'timestamp': time.time()
            }
            
            return result
            
        except Exception as e:
            logger.error(f"推理错误: {e}")
            return None
    
    def inference_worker(self):
        """推理工作线程"""
        logger.info("推理线程启动")
        
        while self.running:
            try:
                # 获取帧
                if not self.frame_queue.empty():
                    frame = self.frame_queue.get(timeout=0.1)
                    
                    # 执行推理
                    result = self.predict(frame)
                    
                    if result:
                        # 将结果放入队列
                        if not self.result_queue.full():
                            self.result_queue.put(result)
                        
                        # 记录日志
                        logger.debug(f"检测到: {result['class_name']} "
                                   f"(置信度: {result['confidence']:.3f}, "
                                   f"风险: {result['risk_level']})")
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"推理线程错误: {e}")
    
    def alert_worker(self):
        """语音提醒工作线程"""
        logger.info("语音提醒线程启动")
        
        last_alert_time = 0
        alert_cooldown = 2.0  # 语音提醒冷却时间(秒)
        
        while self.running:
            try:
                if not self.result_queue.empty():
                    result = self.result_queue.get(timeout=0.1)
                    current_time = time.time()
                    
                    # 检查是否需要语音提醒
                    if (result['risk_level'] in ['medium', 'high'] and 
                        result['confidence'] > 0.7 and
                        current_time - last_alert_time > alert_cooldown):
                        
                        # 触发语音提醒
                        dist_m = self.ultrasonic.get_distance_m()
                        self.voice_alert.alert(
                            result['class_name'], 
                            result['risk_level'],
                            distance_m=dist_m
                        )
                        last_alert_time = current_time
                        
                        logger.info(f"语音提醒: {result['class_name']} "
                                  f"(风险等级: {result['risk_level']})")
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"语音提醒线程错误: {e}")
    
    def calculate_fps(self):
        """计算FPS"""
        self.fps_counter += 1
        current_time = time.time()
        
        if current_time - self.last_fps_time >= 1.0:
            fps = self.fps_counter / (current_time - self.last_fps_time)
            avg_inference_time = np.mean(self.inference_times) if self.inference_times else 0
            
            logger.info(f"FPS: {fps:.1f}, 平均推理时间: {avg_inference_time*1000:.1f}ms")
            
            self.fps_counter = 0
            self.last_fps_time = current_time
    
    def run(self, show_preview: bool = False):
        """运行主循环"""
        logger.info("开始运行障碍检测系统")
        
        self.running = True
        
        # 启动工作线程
        inference_thread = threading.Thread(target=self.inference_worker)
        alert_thread = threading.Thread(target=self.alert_worker)
        
        inference_thread.start()
        alert_thread.start()
        
        try:
            # 主循环 - 摄像头捕获
            for frame in self.camera.capture_frames():
                if not self.running:
                    break
                
                # 将帧放入推理队列
                if not self.frame_queue.full():
                    self.frame_queue.put(frame.copy())
                
                # 显示预览(可选)
                if show_preview:
                    cv2.imshow('Obstacle Detection', frame)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                
                # 计算FPS
                self.calculate_fps()
                
        except KeyboardInterrupt:
            logger.info("接收到中断信号，正在停止...")
        
        finally:
            self.stop()
            inference_thread.join()
            alert_thread.join()
            
            if show_preview:
                cv2.destroyAllWindows()
    
    def stop(self):
        """停止系统"""
        logger.info("正在停止障碍检测系统...")
        self.running = False
        self.camera.stop()
        try:
            self.ultrasonic.stop()
        except Exception:
            pass
        
        # 打印性能统计
        if self.inference_times:
            avg_time = np.mean(self.inference_times)
            min_time = np.min(self.inference_times)
            max_time = np.max(self.inference_times)
            
            logger.info(f"性能统计:")
            logger.info(f"  平均推理时间: {avg_time*1000:.1f}ms")
            logger.info(f"  最快推理时间: {min_time*1000:.1f}ms")
            logger.info(f"  最慢推理时间: {max_time*1000:.1f}ms")
            logger.info(f"  理论最大FPS: {1/avg_time:.1f}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='MobileViT 树莓派障碍检测')
    parser.add_argument('--model', type=str, required=True,
                       help='模型权重文件路径')
    parser.add_argument('--device', type=str, default='cpu',
                       choices=['cpu', 'cuda'],
                       help='推理设备')
    parser.add_argument('--preview', action='store_true',
                       help='显示摄像头预览')
    
    args = parser.parse_args()
    
    # 检查模型文件是否存在
    if not Path(args.model).exists():
        logger.error(f"模型文件不存在: {args.model}")
        return
    
    try:
        # 创建检测器
        detector = ObstacleDetector(args.model, args.device)
        
        # 运行检测
        detector.run(show_preview=args.preview)
        
    except Exception as e:
        logger.error(f"系统运行错误: {e}")
        raise

if __name__ == "__main__":
    main()